#pragma once

#include "../base/types.h"
#include "../base/logger.h"
#include "../base/bits.h"
#include "../base/list.h"
#include <cstddef>
#include <memory>
#include <functional>
#include <vector>

OPEN_JLIB_NS


template<typename Key, typename Val>
struct ordered_map {

  typedef std::pair<Key, Val> Pair;
  typedef Key key_type;
  typedef Val mapped_type;
  typedef Pair value_type;

private:
  constexpr static size_t DefaultBucketSize = 16;

  struct Node {
    Pair pair;
    List list;
    Node* next;
    size_t hash;

    Node (List* list, size_t hash, const Key& key) noexcept : pair(key), list(list->prev(), list), hash(hash) {}

    Node (List* list, size_t hash, Pair&& pair) noexcept : pair(std::move(pair)), list(list->prev(), list), hash(hash) {}

    Node (List* list, size_t hash, const Pair& pair) noexcept : pair(pair), list(list->prev(), list), hash(hash) {}
  };

  typedef Node node_type;

  struct Bucket {
    Node* head;

    bool empty () noexcept { return head == nullptr; }

    void append (Node* node) noexcept { node->next = head; head = node; }

    void remove (Node* node) noexcept {
      debug_assert(head, "bucket is empty");
      if (head == node) {
        head = head->next;
        delete node;
        return;
      }
      for (auto iter = head; iter; iter = iter->next) {
        if (iter->next == node) {
          iter->next = node->next;
          delete node;
          return;
        }
      }
      debug_assert(false, "node not exists");
    }

    Node* find (const Key& key) const noexcept {
      for (auto node = head; node; node = node->next)
        if (node->pair.first == key)
          return node;
      return nullptr;
    }

    void dump (std::vector<Bucket>& buckets) noexcept {
      auto mask = buckets.size() - 1;
      for (auto node = head; node;) {
        auto next = node->next;
        buckets[node->hash & mask].append(node);
        node = next;
      }
      head = nullptr;
    }

    void clear () noexcept {
      for (auto node = head; node;) {
        auto next = node->next;
        delete node;
        node = next;
      }
      head = nullptr;
    }

    Bucket () noexcept : head(nullptr) {}

    ~Bucket () noexcept { clear(); }
  };

public:
  struct const_iterator {
    bool operator!= (const const_iterator& other) const noexcept { return _iter != other._iter; }
    bool operator== (const const_iterator& other) const noexcept { return _iter == other._iter; }
    const_iterator& operator++ () noexcept { _iter = _iter->next(); return *this; }
    const_iterator& operator-- () noexcept { _iter = _iter->prev(); return *this; }
    const_iterator operator++ (int) noexcept { auto old = _iter; _iter = old->next(); return old; }
    const_iterator operator-- (int) noexcept { auto old = _iter; _iter = old->prev(); return old; }
    const Pair& operator* () noexcept { return _iter->cast<&Node::list>()->pair; }
    const Pair* operator-> () noexcept { return &_iter->cast<&Node::list>()->pair; }
  private:
    Node* node () noexcept { return _iter->cast<&Node::list>(); }
    const_iterator (const List* iter) noexcept : _iter(iter) {}
    const List* _iter;
  };

  struct iterator {
    bool operator!= (const iterator& other) const noexcept { return _iter != other._iter; }
    bool operator== (const iterator& other) const noexcept { return _iter == other._iter; }
    iterator &operator++ () noexcept { _iter = _iter->next(); return *this; }
    iterator &operator-- () noexcept { _iter = _iter->prev(); return *this; }
    iterator operator++ (int) noexcept { auto old = _iter; _iter = old->next(); return old; }
    iterator operator-- (int) noexcept { auto old = _iter; _iter = old->prev(); return old; }
    Pair& operator* () noexcept { return _iter->cast<&Node::list>()->pair; }
    Pair* operator-> () noexcept { return &_iter->cast<&Node::list>()->pair; }
  private:
    Node* node () noexcept { return _iter->cast<&Node::list>(); }
    iterator (List* iter) noexcept : _iter(iter) {}
    iterator (const_iterator iter) noexcept : _iter((List*)iter._iter) {}
    List* _iter;
  };

private:
  static size_t __bucket_size (size_t hint) noexcept {
    auto size = Bits<size_t>::highest(hint) << 1;
    return size < DefaultBucketSize ? DefaultBucketSize : size << 1;
  }

  void _rehash (size_t size) noexcept {
    std::vector<Bucket> buckets(size);
    for (auto& bucket : _buckets)
      if (!bucket.empty())
        bucket.dump(buckets);
    _buckets = std::move(buckets);
  }

  template<typename... Args>
  std::pair<iterator, bool> _insert (Args&&... args) noexcept {
    Pair pair(std::forward<Args>(args)...);
    auto size = _buckets.size();
    if (_size >= size)
      _rehash(size <<= 1);
    auto hash = std::hash<Key>()(pair.first);
    auto mask = size - 1;
    auto bucket = _buckets[hash & mask];
    auto node = bucket.find(pair.first);
    if (!node) {
      node = new Node(&_chain, hash, pair);
      bucket.append(node);
      ++_size;
      return {&node->list, true};
    }
    return {&node->list, false};
  }

  template<typename... Args>
  std::pair<iterator, bool> _insert_or_assign (Args&&... args) noexcept {
    Pair pair(std::forward<Args>(args)...);
    auto size = _buckets.size();
    if (_size >= size)
      _rehash(size <<= 1);
    auto hash = std::hash<Key>()(pair.first);
    auto mask = size - 1;
    auto bucket = _buckets[hash & mask];
    auto node = bucket.find(pair.first);
    if (!node) {
      node = new Node(&_chain, hash, pair);
      bucket.append(node);
      ++_size;
      return {&node->list, true};
    }
    node->pair.second = std::move(pair.second);
    return {&node->list, false};
  }

  std::pair<iterator, bool> _get_or_create (const Key& key) noexcept {
    auto size = _buckets.size();
    if (_size >= size)
      _rehash(size <<= 1);
    auto hash = std::hash<Key>()(key);
    auto mask = size - 1;
    auto bucket = _buckets[hash & mask];
    auto node = bucket.find(key);
    if (!node) {
      node = new Node(&_chain, hash, key);
      bucket.append(node);
      ++_size;
      return {&node->list, false};
    }
    return {&node->list, true};
  }

  void _assign (const ordered_map& other) noexcept {
    auto& dst = _buckets;
    auto& src = other._buckets;
    dst.assign(src.size(), Bucket());
    _chain.clear();
    for (auto i = 0; i < src.size(); ++i)
      for (auto node = src[i].head; node; node = node->next)
        dst[i].append(new Node(&_chain, node->hash, node->pair));
    _size = other._size;
  }

public:
  bool empty () const noexcept { return _size == 0; }
  size_t size () const noexcept { return _size; }

  iterator begin () noexcept { return iterator(_chain.next()); }
  iterator rbegin () noexcept { return iterator(_chain.prev()); }
  iterator end () noexcept { return iterator(&_chain); }

  const_iterator begin () const noexcept { return const_iterator(_chain.next()); }
  const_iterator rbegin () const noexcept { return const_iterator(_chain.prev()); }
  const_iterator end () const noexcept { return const_iterator(&_chain); }

  void clear () noexcept {
    for (auto& bucket : _buckets)
      bucket.clear();
    _chain.clear();
    _size = 0;
  }

  std::pair<iterator, bool> insert (const Pair& pair) noexcept {
    return _insert(pair);
  }

  std::pair<iterator, bool> insert (Pair&& pair) noexcept {
    return _insert(std::forward<Pair>(pair));
  }

  template<typename P>
  std::pair<iterator, bool> insert (P&& value) noexcept {
    return _insert(std::forward<P>(value));
  }

  template<typename It>
  void insert (It first, It last) noexcept {
    for (; first != last; ++first)
      _insert(*first);
  }

  void insert (std::initializer_list<Pair> list) noexcept {
    for (auto& pair : list)
      _insert(pair);
  }

  template<typename M>
  std::pair<iterator, bool> insert_or_assign (const Key& key, M&& obj) noexcept {
    return _insert_or_assign(key, std::forward<M>(obj));
  }

  template<typename M>
  std::pair<iterator, bool> insert_or_assign (Key&& key, M&& obj) noexcept {
    auto iter = _insert_or_assign(std::forward<Key>(key), std::forward<M>(obj));
    return std::pair<iterator, bool>(iter, iter != end());
  }

  template<typename M>
  iterator insert_or_assign (const_iterator hint, const Key& key, M&& obj) noexcept {
    if (hint == end() || hint->first != key)
      return _insert_or_assign(key, std::forward<M>(obj)).first;
    iterator iter(hint);
    iter->second = std::forward<M>(obj);
    return iter;
  }

  template<typename M>
  iterator insert_or_assign (const_iterator hint, Key&& key, M&& obj) noexcept {
    if (hint == end() || hint->first != key)
      return _insert_or_assign(std::forward<Key>(key), std::forward<M>(obj)).first;
    iterator iter(hint);
    iter->second = std::forward<M>(obj);
    return iter;
  }

  template<typename... Args>
  std::pair<iterator, bool> emplace (Args&&... args) noexcept {
    return _insert_or_assign(std::forward<Args>(args)...);
  }

  template<typename... Args>
  iterator emplace_hint (const_iterator hint, Args&&... args) noexcept {
    Pair pair(std::forward<Args>(args)...);
    if (hint == end() || hint->first != pair.first)
      return _insert_or_assign(std::move(pair)).first;
    iterator iter(hint);
    iter->second = std::move(pair.second);
    return iter;
  }

  const_iterator find (const Key& key) const noexcept {
    auto size = _buckets.size();
    auto hash = std::hash<Key>()(key);
    auto bucket = _buckets[hash & (size - 1)];
    auto node = bucket.find(key);
    return node ? &node->list : end();
  }

  iterator find (const Key &key) noexcept {
    auto size = _buckets.size();
    auto hash = std::hash<Key>()(key);
    auto bucket = _buckets[hash & (size - 1)];
    auto node = bucket.find(key);
    return node ? &node->list : end();
  }

  iterator erase (const_iterator pos) noexcept {
    debug_assert(pos != end());
    auto iter = ++pos;
    auto node = pos.node();
    auto& bucket = _buckets[node->hash & (_buckets.size() - 1)];
    bucket.remove(node);
    node->list.detach();
    --_size;
    return iter;
  }

  iterator earse (const_iterator first, const_iterator last) noexcept {
    auto iter = last == end() ? last : last++;
    for (; first != last; ++first) {
      auto node = first.node();
      auto& bucket = _buckets[node->hash & (_buckets.size() - 1)];
      bucket.remove(node);
      node->list.detach();
      --_size;
    }
    return iter;
  }

  size_t erase (const Key& key) noexcept {
    auto iter = find(key);
    if (iter != end()) {
      erase(iter);
      return 1;
    }
    return 0;
  }

  void swap (ordered_map& other) noexcept {
    _chain.replace(other._chain);
    _buckets.swap(other._buckets);
    auto size = _size;
    _size = other._size;
    other._size = size;
  }

  Val& operator[] (const Key& key) noexcept {
    return _get_or_create(key).first->second;
  }

  Val& operator[] (Key&& key) noexcept {
    return _get_or_insert(std::forward<Key>(key)).first->second;
  }

  const Val& operator[] (const Key& key) const noexcept {
    return _get_or_create(key).first->second;
  }

  const Val& operator[] (Key&& key) const noexcept {
    return _get_or_insert(std::forward<Key>(key)).first->second;
  }

  explicit ordered_map (size_t size_hint) noexcept : _buckets(__bucket_size(size_hint)), _size(0) {}

  ordered_map () noexcept : ordered_map(DefaultBucketSize) {}

  ordered_map (const ordered_map& other) noexcept {
    _assign(other);
  }

  ordered_map (ordered_map&& other) noexcept {
    _chain.replace(&other.chain);
    _buckets = std::move(other.buckets);
    _size = other._size;
    other._size = 0;
  }

  ordered_map (std::initializer_list<Pair> list) noexcept : _buckets(__bucket_size(list.size())) {
    insert(list);
  }

  ~ordered_map () noexcept {
    clear();
  }

  ordered_map& operator= (const ordered_map& other) noexcept {
    clear();
    _assign(other);
  }

  ordered_map& operator= (ordered_map&& other) noexcept {
    _chain.replace(&other._chain);
    _buckets = std::move(other._buckets);
    _size = other._size;
    other._chain.clear();
    other._size = 0;
  }

  ordered_map& operator= (std::initializer_list<Pair> list) noexcept {
    clear();
    insert(list);
  }

private:
  std::vector<Bucket> _buckets;
  List _chain;
  size_t _size;
};


CLOSE_JLIB_NS
